#!/usr/bin/env python3
# [START program]
"""Vehicles Routing Problem (VRP) for delivering items from any suppliers.

Description: Need to deliver some item X and Y at end nodes (at least 11 X and
13 Y). Several locations provide them and even few provide both.

fleet:
  * vehicles: 2
  * x capacity: 15
  * y capacity: 15
  * start node: 0
  * end node: 1
"""

# [START import]
from ortools.constraint_solver import routing_enums_pb2
from ortools.constraint_solver import pywrapcp

# [END import]


# [START data_model]
def create_data_model():
    """Stores the data for the problem."""
    data = {}
    data["num_vehicles"] = 2
    # [START starts_ends]
    data["starts"] = [0] * data["num_vehicles"]
    data["ends"] = [1] * data["num_vehicles"]
    assert len(data["starts"]) == data["num_vehicles"]
    assert len(data["ends"]) == data["num_vehicles"]
    # [END starts_ends]

    # [START demands_capacities]
    # Need 11 X and 13 Y
    data["providers_x"] = [
        0,  # start
        -11,  # end
        2,  # X supply 1
        2,  # X supply 2
        4,  # X supply 3
        4,  # X supply 4
        4,  # X supply 5
        5,  # X supply 6
        1,  # X/Y supply 1
        2,  # X/Y supply 2
        2,  # X/Y supply 3
        0,  # Y supply 1
        0,  # Y supply 2
        0,  # Y supply 3
        0,  # Y supply 4
        0,  # Y supply 5
        0,  # Y supply 6
    ]
    data["providers_y"] = [
        0,  # start
        -13,  # ends
        0,  # X supply 1
        0,  # X supply 2
        0,  # X supply 3
        0,  # X supply 4
        0,  # X supply 5
        0,  # X supply 6
        3,  # X/Y supply 1
        2,  # X/Y supply 2
        1,  # X/Y supply 3
        3,  # Y supply 1
        3,  # Y supply 2
        3,  # Y supply 3
        3,  # Y supply 4
        3,  # Y supply 5
        5,  # Y supply 6
    ]
    data["vehicle_capacities_x"] = [15] * data["num_vehicles"]
    data["vehicle_capacities_y"] = [15] * data["num_vehicles"]
    assert len(data["vehicle_capacities_x"]) == data["num_vehicles"]
    assert len(data["vehicle_capacities_y"]) == data["num_vehicles"]
    # [END demands_capacities]
    data["distance_matrix"] = [
        [
            0,
            548,
            776,
            696,
            582,
            274,
            502,
            194,
            308,
            194,
            536,
            502,
            388,
            354,
            468,
            776,
            662,
        ],
        [
            548,
            0,
            684,
            308,
            194,
            502,
            730,
            354,
            696,
            742,
            1084,
            594,
            480,
            674,
            1016,
            868,
            1210,
        ],
        [
            776,
            684,
            0,
            992,
            878,
            502,
            274,
            810,
            468,
            742,
            400,
            1278,
            1164,
            1130,
            788,
            1552,
            754,
        ],
        [
            696,
            308,
            992,
            0,
            114,
            650,
            878,
            502,
            844,
            890,
            1232,
            514,
            628,
            822,
            1164,
            560,
            1358,
        ],
        [
            582,
            194,
            878,
            114,
            0,
            536,
            764,
            388,
            730,
            776,
            1118,
            400,
            514,
            708,
            1050,
            674,
            1244,
        ],
        [
            274,
            502,
            502,
            650,
            536,
            0,
            228,
            308,
            194,
            240,
            582,
            776,
            662,
            628,
            514,
            1050,
            708,
        ],
        [
            502,
            730,
            274,
            878,
            764,
            228,
            0,
            536,
            194,
            468,
            354,
            1004,
            890,
            856,
            514,
            1278,
            480,
        ],
        [
            194,
            354,
            810,
            502,
            388,
            308,
            536,
            0,
            342,
            388,
            730,
            468,
            354,
            320,
            662,
            742,
            856,
        ],
        [
            308,
            696,
            468,
            844,
            730,
            194,
            194,
            342,
            0,
            274,
            388,
            810,
            696,
            662,
            320,
            1084,
            514,
        ],
        [
            194,
            742,
            742,
            890,
            776,
            240,
            468,
            388,
            274,
            0,
            342,
            536,
            422,
            388,
            274,
            810,
            468,
        ],
        [
            536,
            1084,
            400,
            1232,
            1118,
            582,
            354,
            730,
            388,
            342,
            0,
            878,
            764,
            730,
            388,
            1152,
            354,
        ],
        [
            502,
            594,
            1278,
            514,
            400,
            776,
            1004,
            468,
            810,
            536,
            878,
            0,
            114,
            308,
            650,
            274,
            844,
        ],
        [
            388,
            480,
            1164,
            628,
            514,
            662,
            890,
            354,
            696,
            422,
            764,
            114,
            0,
            194,
            536,
            388,
            730,
        ],
        [
            354,
            674,
            1130,
            822,
            708,
            628,
            856,
            320,
            662,
            388,
            730,
            308,
            194,
            0,
            342,
            422,
            536,
        ],
        [
            468,
            1016,
            788,
            1164,
            1050,
            514,
            514,
            662,
            320,
            274,
            388,
            650,
            536,
            342,
            0,
            764,
            194,
        ],
        [
            776,
            868,
            1552,
            560,
            674,
            1050,
            1278,
            742,
            1084,
            810,
            1152,
            274,
            388,
            422,
            764,
            0,
            798,
        ],
        [
            662,
            1210,
            754,
            1358,
            1244,
            708,
            480,
            856,
            514,
            468,
            354,
            844,
            730,
            536,
            194,
            798,
            0,
        ],
    ]
    assert len(data["providers_x"]) == len(data["distance_matrix"])
    assert len(data["providers_y"]) == len(data["distance_matrix"])
    return data
    # [END data_model]


# [START solution_printer]
def print_solution(data, manager, routing, assignment):
    """Prints assignment on console."""
    print(f"Objective: {assignment.ObjectiveValue()}")
    # Display dropped nodes.
    dropped_nodes = "Dropped nodes:"
    for node in range(routing.Size()):
        if routing.IsStart(node) or routing.IsEnd(node):
            continue
        if assignment.Value(routing.NextVar(node)) == node:
            dropped_nodes += f" {manager.IndexToNode(node)}"
    print(dropped_nodes)
    # Display routes
    total_distance = 0
    total_load_x = 0
    total_load_y = 0
    for vehicle_id in range(manager.GetNumberOfVehicles()):
        if not routing.IsVehicleUsed(assignment, vehicle_id):
            continue
        index = routing.Start(vehicle_id)
        plan_output = f"Route for vehicle {vehicle_id}:\n"
        route_distance = 0
        route_load_x = 0
        route_load_y = 0
        while not routing.IsEnd(index):
            node_index = manager.IndexToNode(index)
            route_load_x += data["providers_x"][node_index]
            route_load_y += data["providers_y"][node_index]
            plan_output += f" {node_index} Load(X:{route_load_x}, Y:{route_load_y}) -> "
            previous_index = index
            previous_node_index = node_index
            index = assignment.Value(routing.NextVar(index))
            node_index = manager.IndexToNode(index)
            # route_distance += routing.GetArcCostForVehicle(previous_index, index, vehicle_id)
            route_distance += data["distance_matrix"][previous_node_index][node_index]
        node_index = manager.IndexToNode(index)
        plan_output += f" {node_index} Load({route_load_x}, {route_load_y})\n"
        plan_output += f"Distance of the route: {route_distance}m\n"
        plan_output += f"Load of the route: X:{route_load_x}, Y:{route_load_y}\n"
        print(plan_output)
        total_distance += route_distance
        total_load_x += route_load_x
        total_load_y += route_load_y
    print(f"Total Distance of all routes: {total_distance}m")
    print(f"Total load of all routes: X:{total_load_x}, Y:{total_load_y}")
    # [END solution_printer]


def main():
    """Entry point of the program."""
    # Instantiate the data problem.
    # [START data]
    data = create_data_model()
    # [END data]

    # Create the routing index manager.
    # [START index_manager]
    manager = pywrapcp.RoutingIndexManager(
        len(data["distance_matrix"]),
        data["num_vehicles"],
        data["starts"],
        data["ends"],
    )
    # [END index_manager]

    # Create Routing Model.
    # [START routing_model]
    routing = pywrapcp.RoutingModel(manager)

    # [END routing_model]

    # Create and register a transit callback.
    # [START transit_callback]
    def distance_callback(from_index, to_index):
        """Returns the distance between the two nodes."""
        # Convert from routing variable Index to distance matrix NodeIndex.
        from_node = manager.IndexToNode(from_index)
        to_node = manager.IndexToNode(to_index)
        return data["distance_matrix"][from_node][to_node]

    transit_callback_index = routing.RegisterTransitCallback(distance_callback)
    # [END transit_callback]

    # Define cost of each arc.
    # [START arc_cost]
    routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index)
    # [END arc_cost]

    # Add Distance constraint.
    # [START distance_constraint]
    dimension_name = "Distance"
    routing.AddDimension(
        transit_callback_index,
        0,  # no slack
        2000,  # vehicle maximum travel distance
        True,  # start cumul to zero
        dimension_name,
    )
    distance_dimension = routing.GetDimensionOrDie(dimension_name)
    # Minimize the longest road
    distance_dimension.SetGlobalSpanCostCoefficient(100)

    # [END distance_constraint]

    # Add Capacity constraint.
    # [START capacity_constraint]
    def demand_callback_x(from_index):
        """Returns the demand of the node."""
        # Convert from routing variable Index to demands NodeIndex.
        from_node = manager.IndexToNode(from_index)
        return data["providers_x"][from_node]

    demand_callback_x_index = routing.RegisterUnaryTransitCallback(demand_callback_x)
    routing.AddDimensionWithVehicleCapacity(
        demand_callback_x_index,
        0,  # null capacity slack
        data["vehicle_capacities_x"],  # vehicle maximum capacities
        True,  # start cumul to zero
        "Load_x",
    )

    def demand_callback_y(from_index):
        """Returns the demand of the node."""
        # Convert from routing variable Index to demands NodeIndex.
        from_node = manager.IndexToNode(from_index)
        return data["providers_y"][from_node]

    demand_callback_y_index = routing.RegisterUnaryTransitCallback(demand_callback_y)
    routing.AddDimensionWithVehicleCapacity(
        demand_callback_y_index,
        0,  # null capacity slack
        data["vehicle_capacities_y"],  # vehicle maximum capacities
        True,  # start cumul to zero
        "Load_y",
    )
    # [END capacity_constraint]

    # Add constraint at end
    solver = routing.solver()
    load_x_dim = routing.GetDimensionOrDie("Load_x")
    load_y_dim = routing.GetDimensionOrDie("Load_y")
    ends = []
    for v in range(manager.GetNumberOfVehicles()):
        ends.append(routing.End(v))

    node_end = data["ends"][0]
    solver.Add(
        solver.Sum([load_x_dim.CumulVar(l) for l in ends])
        >= -data["providers_x"][node_end]
    )
    solver.Add(
        solver.Sum([load_y_dim.CumulVar(l) for l in ends])
        >= -data["providers_y"][node_end]
    )
    # solver.Add(load_y_dim.CumulVar(end) >= -data['providers_y'][node_end])

    # Allow to freely drop any nodes.
    penalty = 0
    for node in range(0, len(data["distance_matrix"])):
        if node not in data["starts"] and node not in data["ends"]:
            routing.AddDisjunction([manager.NodeToIndex(node)], penalty)

    # Setting first solution heuristic.
    # [START parameters]
    search_parameters = pywrapcp.DefaultRoutingSearchParameters()
    search_parameters.first_solution_strategy = (
        routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC
    )
    search_parameters.local_search_metaheuristic = (
        routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH
    )
    # Sets a time limit; default is 100 milliseconds.
    # search_parameters.log_search = True
    search_parameters.time_limit.FromSeconds(1)
    # [END parameters]

    # Solve the problem.
    # [START solve]
    solution = routing.SolveWithParameters(search_parameters)
    # [END solve]

    # Print solution on console.
    # [START print_solution]
    if solution:
        print_solution(data, manager, routing, solution)
    else:
        print("no solution found !")
    # [END print_solution]


if __name__ == "__main__":
    main()
# [END program]
